import os
import sys
import django
import jieba
import jieba.analyse
from collections import Counter
import re

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'knowledge_site.settings')
django.setup()

from core.models import Document, Tag
from django.utils.text import slugify

def clean_keyword(keyword):
    """清理关键词，移除特殊字符和标点符号"""
    # 移除标点符号和特殊字符
    keyword = re.sub(r'[^\w\s]', '', keyword)
    # 移除空白字符
    keyword = keyword.strip()
    return keyword

def generate_unique_slug(base_slug, existing_slugs, counter=1):
    """生成唯一的slug"""
    if not base_slug:
        base_slug = 'tag'
    new_slug = base_slug
    while new_slug in existing_slugs:
        new_slug = f"{base_slug}-{counter}"
        counter += 1
    return new_slug

def extract_keywords(text, top_k=5):
    """从文本中提取关键词"""
    # 确保文本不为空
    if not text or not isinstance(text, str):
        return []
    return jieba.analyse.extract_tags(text, topK=top_k)

def generate_tags():
    """根据文档生成标签"""
    print("开始生成标签...")
    
    # 获取所有文档
    documents = Document.objects.all()
    if not documents.exists():
        print("没有找到任何文档")
        return
    
    # 收集所有关键词
    all_keywords = []
    for doc in documents:
        # 从标题和内容中提取关键词
        title_keywords = extract_keywords(doc.title, top_k=3)
        content_keywords = extract_keywords(doc.content, top_k=5)
        all_keywords.extend(title_keywords)
        all_keywords.extend(content_keywords)
    
    # 清理和统计关键词
    cleaned_keywords = [clean_keyword(k) for k in all_keywords if clean_keyword(k)]
    keyword_counter = Counter(cleaned_keywords)
    
    # 选择出现频率最高的关键词作为标签
    top_keywords = keyword_counter.most_common(20)
    
    # 获取现有的slugs
    existing_slugs = set(Tag.objects.values_list('slug', flat=True))
    
    # 创建标签
    created_count = 0
    for keyword, count in top_keywords:
        if len(keyword) >= 2:  # 只使用长度大于等于2的关键词
            # 生成基础slug
            base_slug = slugify(keyword)
            if not base_slug:
                continue
                
            # 确保slug唯一
            slug = generate_unique_slug(base_slug, existing_slugs)
            existing_slugs.add(slug)
            
            try:
                tag, created = Tag.objects.get_or_create(
                    name=keyword,
                    defaults={'slug': slug}
                )
                if created:
                    created_count += 1
                    print(f"创建标签: {keyword} (slug: {slug})")
            except Exception as e:
                print(f"创建标签 '{keyword}' 时出错: {str(e)}")
                continue
    
    print(f"\n完成! 共创建了 {created_count} 个新标签")
    
    # 为文档添加标签
    print("\n开始为文档添加标签...")
    for doc in documents:
        try:
            # 获取文档的关键词
            doc_keywords = set(extract_keywords(doc.title, top_k=3) + 
                             extract_keywords(doc.content, top_k=5))
            
            # 清理关键词
            doc_keywords = {clean_keyword(k) for k in doc_keywords if clean_keyword(k)}
            
            # 查找匹配的标签
            tags = Tag.objects.filter(name__in=doc_keywords)
            if tags.exists():
                doc.tags.add(*tags)
                print(f"为文档 '{doc.title}' 添加了 {tags.count()} 个标签")
        except Exception as e:
            print(f"为文档 '{doc.title}' 添加标签时出错: {str(e)}")
            continue
    
    print("\n标签生成和分配完成!")

if __name__ == '__main__':
    generate_tags() 